from flask import Flask, render_template, request, jsonify, Response
import cv2
import numpy as np
from ultralytics import YOLO
import pyautogui
import threading
import time
import os
from datetime import datetime
import json

app = Flask(__name__)

# 全局变量
model = None
is_detecting = False
detected_objects = []
key_press_log = []
selected_classes = []
trigger_key = 'space'
detection_area = {'x': 0, 'y': 0, 'width': 1920, 'height': 1080}

# 安全配置
pyautogui.FAILSAFE = True
pyautogui.PAUSE = 0.1

class YOLODetector:
    def __init__(self):
        self.model = None
        self.is_running = False
        self.thread = None
        
    def load_model(self, model_path):
        try:
            self.model = YOLO(model_path)
            return True, "模型加载成功"
        except Exception as e:
            return False, f"模型加载失败: {str(e)}"
    
    def start_detection(self, classes, key, area):
        global selected_classes, trigger_key, detection_area
        selected_classes = classes
        trigger_key = key
        detection_area = area
        
        if not self.is_running:
            self.is_running = True
            self.thread = threading.Thread(target=self._detect_loop)
            self.thread.daemon = True
            self.thread.start()
            return True, "检测已启动"
        return False, "检测已在运行中"
    
    def stop_detection(self):
        self.is_running = False
        if self.thread:
            self.thread.join()
        return True, "检测已停止"
    
    def _detect_loop(self):
        global detected_objects, key_press_log
        
        while self.is_running:
            try:
                # 截取屏幕
                screenshot = pyautogui.screenshot(region=(
                    detection_area['x'], 
                    detection_area['y'], 
                    detection_area['width'], 
                    detection_area['height']
                ))
                
                # 转换为OpenCV格式
                frame = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)
                
                # YOLO检测
                if self.model:
                    results = self.model(frame)
                    
                    # 获取检测结果
                    detected_objects = []
                    for r in results:
                        boxes = r.boxes
                        if boxes is not None:
                            for box in boxes:
                                cls = int(box.cls[0])
                                conf = float(box.conf[0])
                                
                                # 检查是否在选择的类别中
                                if cls in selected_classes and conf > 0.5:
                                    x1, y1, x2, y2 = box.xyxy[0]
                                    detected_objects.append({
                                        'class': cls,
                                        'confidence': conf,
                                        'bbox': [int(x1), int(y1), int(x2), int(y2)]
                                    })
                                    
                                    # 触发按键
                                    self._trigger_key_press(cls)
                
                time.sleep(0.1)  # 控制检测频率
                
            except Exception as e:
                print(f"检测循环错误: {str(e)}")
                time.sleep(1)
    
    def _trigger_key_press(self, detected_class):
        global key_press_log
        
        try:
            # 模拟按键按下
            pyautogui.press(trigger_key)
            
            # 记录日志
            log_entry = {
                'timestamp': datetime.now().strftime('%H:%M:%S'),
                'class': detected_class,
                'key': trigger_key
            }
            key_press_log.append(log_entry)
            
            # 保持最近5条记录
            if len(key_press_log) > 5:
                key_press_log.pop(0)
                
        except Exception as e:
            print(f"按键触发错误: {str(e)}")

detector = YOLODetector()

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/load_model', methods=['POST'])
def load_model():
    if 'model' not in request.files:
        return jsonify({'success': False, 'message': '没有选择模型文件'})
    
    file = request.files['model']
    if file.filename == '':
        return jsonify({'success': False, 'message': '没有选择模型文件'})
    
    if file and file.filename.endswith('.pt'):
        # 保存模型文件
        model_path = os.path.join('models', file.filename)
        os.makedirs('models', exist_ok=True)
        file.save(model_path)
        
        # 加载模型
        success, message = detector.load_model(model_path)
        return jsonify({'success': success, 'message': message})
    
    return jsonify({'success': False, 'message': '请上传.pt格式的YOLO模型文件'})

@app.route('/get_classes', methods=['GET'])
def get_classes():
    if detector.model:
        try:
            classes = detector.model.names
            return jsonify({'success': True, 'classes': classes})
        except Exception as e:
            return jsonify({'success': False, 'message': str(e)})
    return jsonify({'success': False, 'message': '模型未加载'})

@app.route('/start_detection', methods=['POST'])
def start_detection():
    data = request.json
    classes = data.get('classes', [])
    key = data.get('key', 'space')
    area = data.get('area', {'x': 0, 'y': 0, 'width': 1920, 'height': 1080})
    
    success, message = detector.start_detection(classes, key, area)
    return jsonify({'success': success, 'message': message})

@app.route('/stop_detection', methods=['POST'])
def stop_detection():
    success, message = detector.stop_detection()
    return jsonify({'success': success, 'message': message})

@app.route('/get_status', methods=['GET'])
def get_status():
    return jsonify({
        'is_detecting': detector.is_running,
        'detected_objects': detected_objects,
        'key_press_log': key_press_log
    })

@app.route('/video_feed')
def video_feed():
    def generate():
        while True:
            try:
                # 截取屏幕
                screenshot = pyautogui.screenshot(region=(
                    detection_area['x'], 
                    detection_area['y'], 
                    detection_area['width'], 
                    detection_area['height']
                ))
                
                # 转换为OpenCV格式
                frame = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)
                
                # 绘制检测框
                for obj in detected_objects:
                    x1, y1, x2, y2 = obj['bbox']
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    label = f"Class: {obj['class']}, Conf: {obj['confidence']:.2f}"
                    cv2.putText(frame, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
                
                # 编码为JPEG
                ret, buffer = cv2.imencode('.jpg', frame)
                frame = buffer.tobytes()
                
                yield (b'--frame\r\n'
                       b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
                
            except Exception as e:
                print(f"视频流错误: {str(e)}")
                time.sleep(0.1)
    
    return Response(generate(), mimetype='multipart/x-mixed-replace; boundary=frame')

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)
